#!/usr/bin/env python
import numpy as np
from argparse import ArgumentParser

'''
Script used to randomly generate a gazebo model for the RobotX 2018 "Avoid Obstacles" challenge
'''


def generate_obstacle_course(size, number, min_distance, ul=np.array([0, 0])):
    '''
    Randomly generate a set of random 2D points inside a box subject to the constraint
    that each is greater than a minimum distance apart.
    Note: the algorithm is naive and will block forever if it cannot be solved

    @param size: a 1x2 numpy array representing the size of the box as (X, Y) length
    @param number: the number of points to generate as an integer
    @param min_distance: the minimum euclidian distance two points must be apart to be included in the result
    @return: a number x 2 numpy array of the randomly generated points
    '''
    generated = []
    while len(generated) != number:
        new = np.random.rand(2) * size + ul
        for g in generated:
            if np.linalg.norm(new - g) < min_distance:
                continue
        generated.append(new)
    return np.array(generated)


'''
Python str.format Template for a buoy sub-model, parameterized with the buoys position,
the name to give it in gazebo, and the gazebo model name
'''
TEMPLATE = """
    <include>
      <name>{name}</name>
      <pose>{pos[0]} {pos[1]} {pos[2]} 0 0 0 </pose>
      <uri>model://{model}</uri>
    </include>
"""

# The header to start the model file before the buoy sub-models
HEADER = """<?xml version="1.0"?>
<!-- automaticly generated by the 'generate_avoid_obstacls_buoys' script -->
<sdf version="1.6">
  <model name="robotx_obstacle_course">"""

# The footer to end the model file after the buoy sub-models
FOOTER = """
  </model>
</sdf>"""


def get_models_xml(arr, model):
    '''
    Returns an SDF xml for a set of sub-models
    @param arr: a Nx2 array containing the 2D positions (X, Y) or the models
    @param model: a string for the name of the gazebo model
    @return: a string containing the sdf xml
    '''
    ret = ""
    for idx, pos in enumerate(arr):
        ret += TEMPLATE.format(pos=(pos[0], pos[1], 0.), model=model, name=model + "_" + str(idx))
    return ret

if __name__ == '__main__':
    # Create a command line argument parser
    parser = ArgumentParser(description="Generates an SDF model for the RobotX 2018 obstacle course challenge.")
    parser.add_argument('-s', '--seed', type=int, default=None,
                        help='Seed to use for random buoy placement. Can be used to produce the same model each time.')
    parser.add_argument('--a3', type=int, default=5,
                        help="Number of a3 buoys to include in challenge")
    parser.add_argument('--a5', type=int, default=5,
                        help="Number of a5 buoys to include in challenge")
    parser.add_argument('--a7', type=int, default=5,
                        help="Number of a7 buoys to include in challenge")
    parser.add_argument('--surmark46104', type=int, default=5,
                        help="Number of surmark46104 buoys to include in challenge")
    parser.add_argument('--N', type=int, default=0,
                        help="Generates N markers of each type")
    parser.add_argument('--L', type=float, default=40,
                        help="Size of box used for generation (LxL)")
    args = parser.parse_args()

    if args.seed is not None:
        np.random.seed(args.seed)

    # Get total nubmer of buoys to generate
    total = args.a3 + args.a5 + args.a7 + args.surmark46104

    # Randomly generate the course
    arr = generate_obstacle_course(np.array([args.L, args.L]), total, 5.0, ul=np.array([-args.L/2.0, -args.L/2.0]))

    # Get the model xml for the 3 buoy types using their randomly generated positions
    if args.a3 > 0:
        a3_xml = get_models_xml(arr[0: args.a3, :], "polyform_a3")
    else:
        a3_xml = ""
    if args.a5 > 0:
        a5_xml = get_models_xml(arr[args.a3: args.a3 + args.a5, :], "polyform_a5")
    else:
        a5_xml = ""
    if args.a7 > 0:
        a7_xml = get_models_xml(arr[args.a3 + args.a5: args.a3 + args.a5 + args.a7, :], "polyform_a7")
    else:
        a7_xml = ""
    if args.surmark46104 > 0:
        surmark46104_xml = get_models_xml(arr[args.a3 + args.a5 + args.a7: total, :], "surmark46104")
    else:
        surmark46104_xml = ""


    # Generate a separate course for the N markers  
    N_xml = ""
    if args.N > 0:
        objs = ["surmark46104","surmark950400","surmark950410","black_totem",
                "blue_totem","green_totem","red_totem"]
        for obj in objs:
            arrN = generate_obstacle_course(np.array([args.L, args.L]), args.N, 5.0, ul=np.array([-args.L/2.0, -args.L/2.0])) 
            N_xml += get_models_xml(arrN,obj)

    # Assemble full sdf
    print (HEADER + a3_xml + a5_xml + a7_xml + surmark46104_xml + N_xml + FOOTER)

